import os
import glob
from pathlib import Path

import fitz
import openpyxl

# 这里把书签中可能包含的特殊符号列出来，输出的文件名替换为_，避免创建文件出错
specials = ['?', '*', '/', '\\']


def save_part_as(doc, file_folder, file_name, title, start, end=None):
    for special in specials:
        title = title.replace(special, '_')
    des_file = f"{file_folder}/{file_name}-{start}-{end if end else 'end'}-{title}.pdf"
    # print(des_file)
    if os.path.isfile(des_file):
        return
    doc2 = fitz.open()
    if end:
        doc2.insert_pdf(doc, from_page=start - 1, to_page=end - 2)
    else:
        doc2.insert_pdf(doc, from_page=start - 1)
    print(des_file)
    doc2.save(des_file)


def extract_one_by_pdf(pdf_file, file_folder):
    file_name = os.path.basename(file_folder)
    doc = fitz.open(pdf_file)
    toc = doc.get_toc()
    if not toc:
        print(f'未检测到书签，跳过：{pdf_file}')
        return
    page_title_dic = {}
    prefix_list = []
    last_indent = 1
    last_title = ''
    page_prefix_dic = {}
    for index, ol in enumerate(toc):
        indent, title, page = ol
        # print(indent, title, page)

        if indent > last_indent:
            prefix_list.append(last_title)
        elif prefix_list and indent < last_indent:
            prefix_list.pop()
        # print(prefix_list)
        last_title = title
        last_indent = indent
        page_title_dic[page] = title
        page_prefix_dic[page] = '_'.join(prefix_list) + '_'
    page_title_lst = [(page, title) for page, title in page_title_dic.items()]
    for i in range(len(page_title_lst) - 1):
        page, title = page_title_lst[i]
        # print(page, title)  # page~title, page is start
        end = page_title_lst[i + 1][0]
        save_part_as(doc, file_folder, file_name, page_prefix_dic[page] + title, page, end)
    page, title = page_title_lst[-1]
    save_part_as(doc, file_folder, file_name, page_prefix_dic[page] + title, page)


def extract_bookmarks(pdfs_folder, output_folder):
    src_path = Path(pdfs_folder).absolute().as_posix()
    out_path = Path(output_folder).absolute().as_posix()
    for file in glob.glob(rf'{pdfs_folder}\**\*.pdf', recursive=True):
        file_path = Path(file).absolute().as_posix()
        file_folder = file_path.replace(src_path, out_path).replace('.pdf', '')
        print(file_folder)
        os.makedirs(file_folder, exist_ok=True)
        extract_one_by_pdf(file_path, file_folder)


if __name__ == '__main__':
    pdfs_folder = r'D:\dl\pdfs'
    pdfs_folder = r'C:\Users\yang.lei\Downloads\pdfs'
    output_folder = r'C:\Users\yang.lei\Downloads\dest'
    # output_folder = r'D:\dl\dest'
    extract_bookmarks(pdfs_folder, output_folder)
